#!/usr/bin/env python3
import os
import subprocess
import socket
import json
import sys

WORK_DIR = "/tmp/gpu-burn"
URL = "http://apt.hans.ltd/software/gpu-burn/"
RESULT_PATH = "/tmp/gpu_burn_result.json"
STATUS_URL = "http://data.hans.ltd:5000/test_status"


def run_cmd(cmd, cwd=None, timeout=None):
    """统一运行 shell 命令"""
    try:
        result = subprocess.run(
            cmd,
            cwd=cwd,
            stdout=subprocess.PIPE,
            stderr=subprocess.STDOUT,
            timeout=timeout,
            check=True
        )
        return result.stdout.decode("utf-8", errors="ignore")
    except subprocess.CalledProcessError as e:
        return e.stdout.decode("utf-8", errors="ignore")
    except subprocess.TimeoutExpired:
        return f"TIMEOUT after {timeout}s"
    except Exception as e:
        return f"ERROR: {e}"


def download_repo():
    """用 wget 下载目录平铺到 /tmp/gpu-burn"""
    if os.path.exists(WORK_DIR):
        run_cmd(["rm", "-rf", WORK_DIR])
    os.makedirs(WORK_DIR, exist_ok=True)
    run_cmd([
        "wget",
        "-r", "-nH", "--no-parent", "--cut-dirs=2",
        "-P", WORK_DIR,
        URL
    ])


def build_repo():
    """在 WORK_DIR 下编译 GPU burn"""
    makefile = os.path.join(WORK_DIR, "Makefile")
    if os.path.exists(makefile):
        run_cmd(["make"], cwd=WORK_DIR)
    exe_path = os.path.join(WORK_DIR, "gpu_burn")
    return exe_path if os.path.isfile(exe_path) and os.access(exe_path, os.X_OK) else None


def post_status(status):
    """通过 curl 发送 POST"""
    hostname = socket.gethostname()
    payload = json.dumps({"hostname": hostname, "status": status})
    run_cmd([
        "curl", "-s", "-X", "POST", "-H", "Content-Type: application/json",
        "-d", payload,
        STATUS_URL
    ])


def run_test(exe_path, duration_seconds):
    """运行 GPU burn 测试"""
    output = run_cmd([exe_path, str(duration_seconds)], timeout=duration_seconds + 30)
    result = {
        "hostname": socket.gethostname(),
        "duration_seconds": duration_seconds,
        "output": output
    }
    with open(RESULT_PATH, "w", encoding="utf-8") as f:
        json.dump(result, f)


def main():
    duration_seconds = int(sys.argv[1])

    download_repo()
    exe = build_repo()
    if not exe:
        post_status(2)
        return

    post_status(1)  # 测试开始
    run_test(exe, duration_seconds)
    post_status(2)  # 测试结束


if __name__ == "__main__":
    main()
